from collections import defaultdict, OrderedDict
import datetime
import logging
import traceback
import copy
import re

from coin.exchange.kr_rest.product.product_impl import generate_product_from_str2
from coin.proto.coin_strategy_pb2 import StrategyLog
from coin.proto.coin_order_gateway_pb2 import OrderGatewayLog
from coin.util.queue.tools.kafka_topic import parse_kafka_topic
from coin.base.query_util import query_exchange_rates
from coin.support.pta.logic.pta_stats_calculator import get_futures_multiplier, is_inverse, get_accounting_currency


def _dict_sub(dict1: dict, dict2: dict):
  sub_res = {}
  for key in dict2.keys():
    sub_num = dict1[key] - dict2[key]
    sub_res.update({key: sub_num})
  return sub_res


def _convert_coin(coin):
  sub_res = re.sub(r'U*(10{3,}|[1-9]\d{3,})', '', coin)   # for Phemex U10000SHIB case
  if sub_res == coin:
    multiplier = 1.0
  else:
    multiplier = float(coin.replace(sub_res, ''))
  return sub_res, multiplier  # trailing number GE 1000, e.g. 1000SHIB, 1000LUNC


class BalanceChecker(object):
  def __init__(self, strat_list):
    self._strat_list = strat_list
    self._all_strat_executor_config = defaultdict(dict)
    self._all_strat_balance = defaultdict(dict)
    self._all_strat_position = defaultdict(dict)
    self.balance_extra = []
    self._price_universe = {}

  def _load_price(self, base, quote):
    converted, multiplier = _convert_coin(base)
    price = self._price_universe.get(f'{converted}-{quote}', None)
    price *= multiplier
    return price

  def _query_exchange_rates(self, bases, quote):
    gen_symbol = lambda base, quote: '%s-%s' % (base, quote)
    bases_to_query = []
    for base in bases:
      converted, _ = _convert_coin(base)
      if gen_symbol(converted, quote) not in self._price_universe:
        bases_to_query.append(converted)
    if len(bases_to_query) > 0:
      prices = query_exchange_rates(bases_to_query, quote)
      price_dict = dict(zip(bases_to_query, prices))
      failed = [b for p, b in zip(prices, bases) if p is None]
      prev_day = datetime.datetime.utcnow().date() - datetime.timedelta(days=1)
      if failed:
        history_prices = query_exchange_rates(failed, quote, prev_day)
        price_dict.update(dict(zip(failed, history_prices)))
      for base, price in price_dict.items():
        if price is not None:
          self._price_universe[gen_symbol(base, quote)] = price

    return [self._load_price(base, quote) for base in bases]

  @staticmethod
  def query_exchange_rates_by_dates(base: str, quote: str, dates: list) -> OrderedDict:
    bases_to_query = []
    converted, multiplier = _convert_coin(base)
    rev_dates = dates.copy()
    rev_dates.reverse()
    res = OrderedDict()
    bases_to_query.append(converted)
    if len(bases_to_query) > 0:
      for date in rev_dates:
        price = query_exchange_rates(bases_to_query, quote, date)[0]
        if price is not None:
          res[date] = price * multiplier
        elif price is None and len(res) > 0:
          res.append({date: list(res.items())[-1][1]})
    if len(res) == 0:
      tmp_dates = [
          rev_dates[-1] - datetime.timedelta(days=1),
          datetime.datetime.utcnow().date(),
          datetime.datetime.utcnow().date() - datetime.timedelta(days=1)
      ]
      for date in tmp_dates:
        price = query_exchange_rates(bases_to_query, quote, date)[0]
        if price is None:
          continue
        else:
          res[date] = price * multiplier
          break
    return res


  def _handle_og_log(self, strat_name, pb):
    assert pb.HasField('og_log'), ("log type %s has no og_log" % pb.type)
    og_log = pb.og_log
    if og_log.type == OrderGatewayLog.BALANCE:
      self._update_balance_info(strat_name, og_log)
    elif og_log.type == OrderGatewayLog.POSITION:
      self._update_position_info(strat_name, og_log)

  def _handle_executor_config(self, strat_name, pb):
    if pb.type == StrategyLog.AGG_EXEUCTOR_CONFIG:
      assert pb.HasField('agg_executor_config')
      executor_config = pb.agg_executor_config
    elif pb.type == StrategyLog.PASS_EXECUTOR_CONFIG:
      assert pb.HasField('pass_executor_config')
      executor_config = pb.pass_executor_config
    else:
      raise ValueError('Unknown executor type: %s' % pb.type)
    if not executor_config.symbol:
      return
    account_request = None
    if pb.HasField('account_request'):
      account_request = pb.account_request
    if executor_config.HasField('account_request'):
      account_request = executor_config.account_request
    self._update_executor_config(strat_name, account_request, executor_config)

  def _update_position_info(self, strat_name, og_log):
    assert og_log.HasField('account_request')
    account_request = og_log.account_request
    position = {}
    for each_position in og_log.position.each_position:
      position[each_position.symbol] = each_position
    account_key = (account_request.market_type, account_request.exchange, account_request.owner)
    orig_position = self._all_strat_position[strat_name].get(account_key, {})
    orig_position.update(position)
    self._all_strat_position[strat_name].update({account_key: orig_position})

  def _update_balance_info(self, strat_name, og_log):
    assert og_log.HasField('account_request')
    account_request = og_log.account_request
    strat_balance = og_log.balance
    balance = {
        each_balance.currency: each_balance.total for each_balance in strat_balance.each_balance
    }
    account_key = (account_request.market_type, account_request.exchange, account_request.owner)
    orig_balance = self._all_strat_balance[strat_name].get(account_key, {})
    # sometimes balance doesn't include all coins, but only part of them, so need to
    # use dict update
    orig_balance.update(balance)
    self._all_strat_balance[strat_name].update({account_key: orig_balance})

  def _update_executor_config(self, strat_name, account_request, executor_config):
    account_key = None
    if account_request is None:
      if len(self._all_strat_balance[strat_name]) == 1:
        account_key = list(self._all_strat_balance[strat_name].keys())[0]
    else:
      account_key = (account_request.market_type, account_request.exchange, account_request.owner)
    if account_key is None:
      return
    product = generate_product_from_str2(account_key[0], account_key[1], None,
                                         executor_config.symbol)
    symbol = product.symbol
    symbol_config = {}
    symbol_config['max_pos'] = max(abs(executor_config.max_pos), abs(executor_config.min_pos))
    account_max_config = self._all_strat_executor_config.get(strat_name, {}).get(account_key, {})
    account_max_config[symbol] = symbol_config
    self._all_strat_executor_config[strat_name].update({account_key: account_max_config})

  def calculate_all_extra_balance(self):
    extra_balance = []
    for strat_name in self._strat_list:
      try:
        strategy_extra_balance = self._calculate_strategy_extra_balance(strat_name)
        extra_balance.extend(strategy_extra_balance)
      except Exception as e:
        print('Skipping strategy:', strat_name, type(e), e, traceback.format_exc())
    return extra_balance

  def _calculate_strategy_extra_balance(self, strat_name):
    strat_extra_balance = []
    assert strat_name in self._all_strat_balance, "No balance info"
    for account_key, balance in self._all_strat_balance[strat_name].items():
      market_type, exchange, owner = account_key
      if strat_name not in self._all_strat_executor_config or \
         account_key not in self._all_strat_executor_config[strat_name]:
        raise ValueError("No executor config: %s" % (account_key,))
      max_config = self._all_strat_executor_config[strat_name][account_key]
      if market_type == 'Spot':
        acct_extra_balance = self._calculate_spot_extra_balance(account_key, balance, max_config)
      elif market_type == 'Futures':
        acct_extra_balance = self._calculate_futures_extra_balance(account_key, balance, max_config)
      else:
        raise ValueError('Unsupported market type: %s' % market_type)
      acct_extra_balance['strat_name'] = strat_name
      acct_extra_balance['exchange'] = exchange
      acct_extra_balance['owner'] = owner
      acct_extra_balance['market_type'] = market_type
      acct_extra_balance['config'] = max_config
      strat_extra_balance.append(acct_extra_balance)
    return strat_extra_balance

  def _calculate_spot_extra_balance(self, account_key, balance, max_config):
    market_type, exchange, symbol = account_key
    assert market_type == 'Spot', account_key
    min_value = 0
    account_value = 0
    symbol_price = {}
    quote = None
    for symbol, symbol_config in max_config.items():
      product = generate_product_from_str2(market_type, exchange, None, symbol)
      prod_accounting_currency = get_accounting_currency(market_type, exchange, symbol)
      if quote is not None:
        assert prod_accounting_currency == quote, \
            "Multi accounting currency in spot account: %s, %s vs %s" % (
                account_key, prod_accounting_currency, quote)
      quote = prod_accounting_currency
      symbol = product.symbol
      price = self._query_exchange_rates([product.base.currency], product.quote.currency)[0]
      if price is None:
        logging.info("price not found: %s %s" % (exchange, symbol))
        continue
      symbol_price.update({symbol: price})
      min_value += symbol_config['max_pos'] * price
      base = product.base.currency
      account_value += balance.get(base, 0) * price
    account_value += balance.get(quote, 0)
    account_balance = {quote: account_value}
    min_balance = {quote: min_value}
    extra_balance = {quote: account_value - min_value}
    return {
        'cur_balance': account_balance,
        'min_balance': min_balance,
        'extra_balance': extra_balance,
        'symbol_price': symbol_price,
    }

  def _calculate_futures_extra_balance(self, account_key, balance, max_config):
    leverage = 12
    market_type, exchange, owner = account_key
    assert market_type == 'Futures', account_key
    min_balance_map = {}
    symbol_price = {}
    for symbol, symbol_config in max_config.items():
      product = generate_product_from_str2(market_type, exchange, None, symbol)
      symbol = product.symbol
      base = product.base.currency
      quote = product.quote.currency
      price = self._query_exchange_rates([base], quote)[0]
      if price is None:
        logging.info("price not found: %s %s" % (exchange, symbol))
        continue
      multiplier = get_futures_multiplier(market_type, exchange, symbol)
      accounting_currency = get_accounting_currency('Futures', exchange, symbol)
      if is_inverse(market_type, exchange, symbol):
        min_balance = symbol_config['max_pos'] * multiplier / leverage / price
      else:
        min_balance = symbol_config['max_pos'] * price * multiplier / leverage
      min_balance += min_balance_map.get(accounting_currency, 0)
      min_balance_map.update({accounting_currency: min_balance})
      symbol_price.update({symbol: price})
      symbol_config.update({'leverage': leverage, 'multiplier': multiplier})
      max_config.update({symbol: symbol_config})
    try:
      if exchange == 'Ftx':
        adj_balance = self._get_ftx_futures_margin(balance, list(min_balance_map.keys()))
        extra_balance = _dict_sub(adj_balance, min_balance_map)
      elif exchange == 'Binance' and 'portfolio' in owner:
        balance, min_balance_map = self._get_binance_portfolio_margin(balance, min_balance_map)
        extra_balance = _dict_sub(balance, min_balance_map)
      else:
        extra_balance = _dict_sub(balance, min_balance_map)
    except:
      print(balance, min_balance_map)
      raise
    return {
        'cur_balance': balance,
        'min_balance': min_balance_map,
        'extra_balance': extra_balance,
        'symbol_price': symbol_price,
    }

  def _get_ftx_futures_margin(self, balance, ccy_list):
    assert len(ccy_list) == 1 and ccy_list[0] == 'USD'
    adj_balance = copy.copy(balance)
    margin_list = ('BTC', 'ETH')
    price_list = self._query_exchange_rates(margin_list, 'USD')
    quote_asset = balance.get('USD', 0.0)
    for i in range(len(margin_list)):
      coin_asset = balance.get(margin_list[i], 0.0)
      quote_asset += coin_asset * price_list[i]
      adj_balance['USD'] = quote_asset
    return adj_balance

  def _get_binance_portfolio_margin(self, balance, min_balance_map):
    adj_balance = copy.copy(balance)
    coin_list = list(adj_balance.keys())
    price_list = self._query_exchange_rates(coin_list, 'USD')
    margin_in_usd = sum([adj_balance.get(coin_list[i], 0) * price_list[i]
                        for i in range(len(coin_list))])
    min_coin_list = list(min_balance_map.keys())
    min_price_list = self._query_exchange_rates(min_coin_list, 'USD')
    min_margin_in_usd = sum([min_balance_map.get(min_coin_list[i], 0) *
                            min_price_list[i] for i in range(len(min_coin_list))])
    adj_balance['USD'] = margin_in_usd
    adj_min_balance_map = {'USD': min_margin_in_usd}
    return adj_balance, adj_min_balance_map

  def handle(self, partition, record):
    topic = record.topic
    pb = StrategyLog()
    pb.ParseFromString(record.value)
    strat_name = parse_kafka_topic(topic)[1].strategy_name
    try:
      if pb.type in (StrategyLog.PASS_EXECUTOR_CONFIG, StrategyLog.AGG_EXEUCTOR_CONFIG):
        self._handle_executor_config(strat_name, pb)
      elif pb.type == StrategyLog.OG_LOG:
        self._handle_og_log(strat_name, pb)
      else:
        pass
    except Exception as e:
      logging.info(
          'Fail to handle %s, %s, %s' % (record.timestamp * 1e6, strat_name, pb))
